#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Date    : 2018-08-21 10:57:30
# @Author  : wanglele (18911756410@163.com)
# @Link    : None
# @Version : "Python 3.7"

import os
import sys
import pickle


class URLManager:
    def __init__(self):
        self.new_urls = self.load_progress("new_urls.txt")
        self.old_urls = self.load_progress("old_urls.txt")

    def load_progress(self, path):
        """
        从文件加载进度
        path: 文件路径
        返回一个集合
        """
        print("[+] 从文件加载进度： %s" % path)
        try:
            with open(path, "rb") as f:
                tmp = pickle.load(f)
                return tmp
        except Exception as e:
            print("[!] 无进度文件，创建：%s" % path)
        return set()

    def save_progress(self, path, data):
        """
        保存进度
        path 文件路径
        dara 数据
        """
        with open(path, "wb") as f:
            pickle.dump(data, f)

    def new_urls_size(self):
        """
        获取未爬取url集合大小
        """
        return len(self.new_urls)

    def old_urls_size(self):
        """
        获取已爬取url集合大小
        """
        return len(self.old_urls)

    def add_new_url(self, url):
        """
        将新的url添加到未爬取url集合中
        """
        if url is None:
            return None
        if url not in self.new_urls and url not in self.old_urls:
            self.new_urls.add(url)

    def add_new_urls(self, urls):
        """
        将新的url添加到未爬取url集合中
        """
        if urls is None or len(urls) == 0:
            return None
        for url in urls:
            self.add_new_url(url)

    def get_new_url(self):
        """
        从未爬取的url集合中拿出一个要爬的url
        """
        new_url = self.new_urls.pop()
        self.old_urls.add(new_url)
        return new_url

    def has_new_url(self):
        """
        判断是否有未爬取的url
        """
        return self.new_urls_size() != 0
